import matplotlib.pyplot as plt
import numpy as np
from collections import defaultdict
import time


def array2index(array):
	index = ''
	for n, num in enumerate(array):
		index += str(num)+' '
	return index

# ts1 = [2,1,3,2,1,3,2,1,3,2]
# ts2 = [1,2,3,1,2,3,1,2,3,1]


def symbol(subsequence):
	return array2index(np.argsort(subsequence))
	

def TE(ts1, ts2, l, m, step=1, func=symbol):
	N = len(ts1)
	
	x = []
	y = []
	
	for i in range(N - l * (m - 1)):
		subsequence1 = ts1[i:i + l * (m - 1) + 1:l]
		subsequence2 = ts2[i:i + l * (m - 1) + 1:l]
		k_x = func(subsequence1)
		k_y = func(subsequence2)
		x.append(k_x)
		y.append(k_y)
	
	sequence_x1 = x[step:]
	sequence_x0 = x[:-step]
	sequence_y0 = y[:-step]
	Nx1x0y0 = defaultdict(int)
	Nx1x0 = defaultdict(int)
	Nx0y0 = defaultdict(int)
	Nx0 = defaultdict(int)
	N = len(sequence_x1)
	
	for i in range(N):
		Nx1x0y0[(sequence_x1[i], sequence_x0[i], sequence_y0[i])] += 1
		Nx1x0[(sequence_x1[i], sequence_x0[i])] += 1
		Nx0y0[(sequence_x0[i], sequence_y0[i])] += 1
		Nx0[sequence_x0[i]] += 1
	
	# print(Nx1x0y0, Nx1x0, Nx0y0, Nx0)
	
	px1dx0y0 = defaultdict(int)
	px1dx0 = defaultdict(int)
	px1x0y0 = defaultdict(int)
	
	for (x1, x0, y0) in Nx1x0y0.keys():
		if Nx1x0[(x1, x0)] != 0:
			px1dx0y0[(x1, x0, y0)] = Nx1x0y0[(x1, x0, y0)] / Nx0y0[(x0, y0)]
		if Nx0[x0] != 0:
			px1dx0[(x1, x0)] = Nx1x0[(x1, x0)] / Nx0[x0]
		if Nx0y0[(x0, y0)] != 0:
			px1x0y0[(x1, x0, y0)] = Nx1x0y0[(x1, x0, y0)] / N
	
	# print('N calculation done')
	Tyx = 0
	# print(px1dx0y0, px1x0y0, px1dx0)
	for (x1, x0, y0) in Nx1x0y0.keys():
		if px1dx0[(x1, x0)] != 0 and px1x0y0[(x1, x0, y0)] != 0:
			Tyx += px1x0y0[(x1, x0, y0)] * np.log2(px1dx0y0[(x1, x0, y0)] / px1dx0[(x1, x0)])

	return Tyx


if __name__ == '__main__':
	ts1 = np.loadtxt('../data/pr_NY.csv')  # 降水量 # Y
	ts2 = np.loadtxt('../data/tas_NY.csv')  # 温度 # X
	print('...')
	time_start = time.time()
	# STE_true_case(ts1, ts2, 1, 5, step=1)
	l = 1
	m = 3
	step = 1
	T_series = []
	for m in range(3, 20):
		T = TE(ts2, ts1, l, m, step) - TE(ts1, ts2, l, m, step)
		T_series.append(T)
	time_end = time.time()
	plt.plot(T_series)
	plt.show()

	# print('Cost:', time_end - time_start)
